import torch
from torch import nn
import torch.nn.functional as F
from VAEbase import BaseVAE
from convlstm import ConvLSTMCell

class ResBlock(nn.Module):
    def __init__(self, in_channels, num_hiddens, num_residual_hiddens):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = num_hiddens
        self._block = nn.Sequential(
            nn.ReLU(),
            nn.Conv2d(in_channels=in_channels,
                      out_channels=num_residual_hiddens,
                      kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(num_residual_hiddens),
            nn.ReLU(),
            nn.Conv2d(in_channels=num_residual_hiddens,
                      out_channels=num_hiddens,
                      kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(num_hiddens)
        )
        if in_channels != num_hiddens:
            self.conv_shortcut = nn.Conv2d(in_channels=in_channels,
                                        out_channels=num_hiddens,
                                        kernel_size=3,
                                        stride=1,
                                        padding=1)
        
    def forward(self, x):
        h = x
        if self.in_channels != self.out_channels:
            x = self.conv_shortcut(x)
            h = self._block(h)
            return x + h
        else:
            return x + self._block(x)
    
class ResidualStack(nn.Module):
    def __init__(self, in_channels, num_hiddens, num_residual_layers, num_residual_hiddens):
        super().__init__()
        self._num_residual_layers = num_residual_layers
        res_stack = [ResBlock(in_channels, num_hiddens, num_residual_hiddens)]
        for _ in range(1, self._num_residual_layers):
            res_stack.append(ResBlock(num_hiddens, num_hiddens, num_residual_hiddens))
        self._layers = nn.ModuleList(res_stack)
        
    def forward(self, x):
        for i in range(self._num_residual_layers):
            x = self._layers[i](x)
        return F.relu(x)
    
    
class Encoder(nn.Module):
    def __init__(self, 
                 in_channels, 
                 num_hiddens, 
                 num_residual_layers, 
                 num_residual_hiddens,
                 z_channels):
        super().__init__()
        self.conv1 = nn.Sequential(nn.Conv2d(in_channels=in_channels,
                                              out_channels=num_hiddens//2,
                                              kernel_size=4,
                                              stride=2,
                                              padding=1),
                                    nn.BatchNorm2d(num_hiddens//2),
                                    nn.ReLU())
        self.conv2 = nn.Sequential(nn.Conv2d(in_channels=num_hiddens//2,
                                              out_channels=num_hiddens,
                                              kernel_size=4,
                                              stride=2,
                                              padding=1),
                                    nn.BatchNorm2d(num_hiddens))
        self.residual_stack = ResidualStack(in_channels=num_hiddens,
                                             num_hiddens=z_channels * 2,
                                             num_residual_layers=num_residual_layers,
                                             num_residual_hiddens=num_residual_hiddens)
    
    def forward(self, inputs):
        x = self.conv1(inputs)
        x = self.conv2(x)
        enc_out = self.residual_stack(x)
        return enc_out
        
    
class Decoder(nn.Module):
    def __init__(self, out_channels, num_hiddens, num_residual_layers, num_residual_hiddens):
        super().__init__()
        self._residual_stack = ResidualStack(in_channels=num_hiddens,
                                             num_hiddens=num_hiddens,
                                             num_residual_layers=num_residual_layers,
                                             num_residual_hiddens=num_residual_hiddens)
        self._conv_trans2 = nn.Sequential(nn.ReLU(),
                                          nn.ConvTranspose2d(in_channels=num_hiddens,
                                                             out_channels=num_hiddens//2,
                                                             kernel_size=4,
                                                             stride=2,
                                                             padding=1),
                                          nn.BatchNorm2d(num_hiddens//2),
                                          nn.ReLU())
        self._conv_trans1 = nn.Sequential(nn.ConvTranspose2d(in_channels=num_hiddens//2,
                                                             out_channels=num_hiddens//2,
                                                             kernel_size=4,
                                                             stride=2,
                                                             padding=1),
                                          nn.BatchNorm2d(num_hiddens//2),
                                          nn.ReLU(),
                                          nn.Conv2d(in_channels=num_hiddens//2,
                                                    out_channels=out_channels,
                                                    kernel_size=3,
                                                    stride=1,
                                                    padding=1),
                                          nn.Sigmoid())
        
    def forward(self, inputs):
        x = self._residual_stack(inputs)
        x = self._conv_trans2(x)
        x = self._conv_trans1(x)
        return x
    
class DenseCondVAE(BaseVAE):
    def __init__(self,
                 encoderconfig,
                 decoderconfig,
                 convlstmconfig):
        super().__init__()
        self.convlstm = ConvLSTMCell(**convlstmconfig)
        self.encoder = Encoder(**encoderconfig)
        self.decoder = Decoder(**decoderconfig)
        